
from flask import Flask, render_template, request, send_from_directory
import os
import cv2
from werkzeug.utils import secure_filename
from datetime import datetime
import hashlib
import shutil

current_path = os.getcwd()
SUPPORT_IMAGE_TYPE = ('.png', '.jpg', '.jpeg', '.gif')
SUPPORT_VIDEO_TYPE = ('.mp4', '.avi', '.mov')


app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = r'E:\weibo\视频任务组_20230703_0009'
app.config['THUMBNAIL_FOLDER'] = os.path.join(current_path, "static", "thumbnails")

# 确保目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['THUMBNAIL_FOLDER'], exist_ok=True)

def generate_video_thumbnail(video_path, video_filename):
    """生成视频预览图"""
    cap = cv2.VideoCapture(video_path)
    success, image = cap.read()
    if success:
        thumbnail_name = video_filename + '.jpg'
        thumbnail_path = os.path.join(app.config['THUMBNAIL_FOLDER'], thumbnail_name)
        cv2.imwrite(thumbnail_path, image)
        hash_value = dhash(image)
        return thumbnail_name, hash_value
    return None, None

@app.route('/')
def index():
    delete_directory_contents(app.config['THUMBNAIL_FOLDER'])
    media_files = []
    for filename in os.listdir(app.config['UPLOAD_FOLDER']):
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        if filename.lower().endswith(SUPPORT_IMAGE_TYPE):
            media_files.append({
                'name': filename,
                'type': 'image',
                'thumbnail': f'uploads/{filename}',
                'thumbnail_hash': hashlib.md5(open(filepath, 'rb').read()).hexdigest()
            })
        elif filename.lower().endswith(SUPPORT_VIDEO_TYPE):
            thumbnail, hash_value = generate_video_thumbnail(filepath, filename)
            media_files.append({
                'name': filename,
                'type': 'video',
                'thumbnail': f'thumbnails/{thumbnail}' if thumbnail else None,
                'thumbnail_hash': hash_value
            })
    
    # 按预览图哈希分组
    grouped_media = {}
    for item in media_files:
        if item['thumbnail_hash'] not in grouped_media:
            grouped_media[item['thumbnail_hash']] = []
        grouped_media[item['thumbnail_hash']].append(item)

    grouped_list = []
    depulicate_file_num = 0
    for k, v in grouped_media.items():
        grouped_list.append((k, v))
        depulicate_file_num += len(v) - 1
    
    return render_template('index.html', grouped_media=sorted(grouped_list, key=lambda x: -len(x[1])), depulicate_file_num=depulicate_file_num)

@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

@app.route('/thumbnails/<filename>')
def thumbnail_file(filename):
    return send_from_directory(app.config['THUMBNAIL_FOLDER'], filename)

def dhash(image, hash_size=16):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    resized = cv2.resize(gray, (hash_size+1, hash_size))
    diff = resized[:,1:] > resized[:,:-1]
    return sum([2**i for (i,v) in enumerate(diff.flatten()) if v])

def delete_directory_contents(directory):
    for filename in os.listdir(directory):
        file_path = os.path.join(directory, filename)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            print(f"删除{file_path}失败: {e}")

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=8081, debug=True)
